"""
    @brief:The gas mass fraction is introduced and the grid distribution is used for clustering
    @Editor:CJH
    @Date:2025/3/13
"""
import os

import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.ndimage import gaussian_filter

np.set_printoptions(suppress=True)


class Dataloader:
    def __init__(self, grid_shape):
        self.grid_shape = grid_shape  # 转换的质量分数栅格地图,(y,x)
        self.init_map = np.ones(grid_shape, dtype=np.float32)
        self.noisy_map = np.ones(grid_shape, dtype=np.float32)

    @staticmethod
    def load_csv(csv_path):  # 读取csv表格数据
        assert os.path.exists(csv_path), "Please ensure input the correct path of csvfile"
        df = pd.read_csv(csv_path, index_col=False)
        try:
            df = df.set_index('Node Number')
        except:
            pass
        return df

    def add_Gaussian_noisy(self, sigma: float, map: np.ndarray):
        # 步骤1：随机移位（模拟位置偏移）
        shift_x = np.random.randint(-5, 5)  # x方向随机移位 [-2, 2] 格
        shift_y = np.random.randint(-5, 5)  # y方向随机移位 [-2, 2] 格

        shifted_map = np.roll(map, shift=(shift_y, shift_x), axis=(0, 1))

        # 步骤2：应用高斯滤波，模拟位置偏移
        blurred_map = gaussian_filter(shifted_map, sigma=sigma)

        # 步骤3：添加高斯噪声
        noise_std = 0.01  # 噪声标准差
        noisy_map = blurred_map + np.random.normal(loc=0, scale=noise_std, size=map.shape)

        # 步骤4：确保值非负
        noisy_map = np.clip(noisy_map, 0, None)

        # 步骤5：可选归一化
        noisy_map = noisy_map / np.sum(noisy_map) * np.sum(map)

        # 更新 init_map
        self.noisy_map = noisy_map

    # 将表格数据进行聚类分割,表格数据与地图数据融合
    def merge(self, csv_path):

        # 获取表格数据
        dataframe = self.load_csv(csv_path)
        # 由于方向都是正的,abs
        dataframe.values[:, :] = abs(dataframe.values[:, :])
        # 预先获取列名和最大值
        x_row_name = dataframe.columns[1]  # 'Y [ m ]',由于y坐标分割行
        y_col_name = dataframe.columns[0]  # 'X [ m ]'，由于x坐标分割列
        value_col_name = dataframe.columns[-1]  # 'C2h5oh.Mass Fraction'
        df = dataframe  # 假设 df 是 dataframe 的引用
        max_x = df[x_row_name].max()
        max_y = df[y_col_name].max()

        # 创建网格的分箱边界
        x_bins = np.linspace(0, max_x, self.grid_shape[1] + 1)
        y_bins = np.linspace(0, max_y, self.grid_shape[0] + 1)

        # 将 X 和 Y 列分箱到网格
        df['x_bin'] = pd.cut(df[x_row_name], bins=x_bins, labels=False, include_lowest=True)
        df['y_bin'] = pd.cut(df[y_col_name], bins=y_bins, labels=False, include_lowest=True)

        # 按网格分组并计算均值
        grid_means = df.groupby(['x_bin', 'y_bin'])[value_col_name].mean().unstack(fill_value=0)

        # 将结果填入 init_map 并处理 NaN 和负值
        self.init_map = grid_means.values.T  # 转置以匹配 (x_col, y_row) 顺序
        self.init_map = np.where(np.isnan(self.init_map), 0, self.init_map)  # NaN 转为 0

        if 50 < np.sum(self.grid_shape) < 150:
            self.init_map = cv2.blur(self.init_map, (2, 2))
        elif 150 < np.sum(self.grid_shape):
            self.init_map = cv2.blur(self.init_map, (3, 3))

        self.add_Gaussian_noisy(1, self.init_map)

        # plt.clf()
        # plt.imshow(self.init_map, cmap='jet')  # cmap 指定颜色映射
        # plt.colorbar(label='Value')  # 添加颜色条
        # plt.title('Heatmap Example')  # 设置标题
        # plt.xlabel('X-axis')  # X 轴标签
        # plt.ylabel('Y-axis')  # Y 轴标签
        # plt.gca().invert_yaxis()  # 反转 y 轴
        #
        #
        # plt.figure()
        # plt.imshow(self.noisy_map, cmap='jet')  # cmap 指定颜色映射
        # plt.colorbar(label='Value')  # 添加颜色条
        # plt.title('Heatmap Example:noisy-map')  # 设置标题
        # plt.xlabel('X-axis')  # X 轴标签
        # plt.ylabel('Y-axis')  # Y 轴标签
        # plt.gca().invert_yaxis()  # 反转 y 轴
        # plt.show()

        # # plt.pause(0.001)
        # plt.cla()
        # plt.close()


if __name__ == "__main__":
    dataloader = Dataloader((45, 30))
    dataloader.merge(r"D:\Gas_detector\data\gas_mass_data\export4.csv")
